import dotenv

from openai import OpenAI
from langchain_community.chat_message_histories import FileChatMessageHistory

dotenv.load_dotenv()

# 1. 创建客户端&记忆
client = OpenAI(base_url='https://api.moonshot.cn/v1')
chat_history = FileChatMessageHistory("./memory.txt")

# 2. 循环对话
while True:
    # 3. 获取用户的输入
    query = input("Human: ")

    # 4. 检测用户是否退出对话
    if query == "q":
        exit(0)

    # 5. 发起聊天对话
    print("AI: ", flush=True, end="")
    system_prompt = (
        "你是OpenAI开发的ChatGPT机器人， 可以根据响应的上下文回复用户的信息，上下文里存放的是人类与你对话的信息列表。\n\n"
        f"<context>{chat_history}</context>\n\n"
    )

    response = client.chat.completions.create(
        model="kimi-k2-0711-preview",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": query}
        ],
        stream=True
    )
    ai_content = ""
    for chunk in response:
        content = chunk.choices[0].delta.content
        if content is None:
            break
        ai_content += content
        print(content, flush=True, end="")
    chat_history.add_user_message(query)
    chat_history.add_ai_message(ai_content)
    print("")
